# -*- coding: UTF-8 -*-

from pymongo import MongoClient
from WindPy import *

import datetime
import math

from StockDailyFactorGeneralFunctions import StockDailyFactorGeneralFunctions

class DailyClose:
    def __init__(self):
        self.baseStartDateStr = '2012-01-01'
        self.baseEndDateStr = '2017-05-01'

    def writeBaseStocksCloseFromWindToMongoDB(self):
        baseStartDateStr = self.baseStartDateStr
        baseEndDateStr = self.baseEndDateStr

        windFactorNameStr = 'close'
        colNameStr = 'daily_close'
        dbUpdateDateFiledStr = 'stock_close_update_date'
        dbUpdateDateFieldAliasStr = ' daily close '

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.writeBaseDataToMongoDB(baseStartDateStr, baseEndDateStr, windFactorNameStr, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr)

    def writeBaseAdjFactorFromWindToMongoDB(self):
        baseStartDateStr = self.baseStartDateStr
        baseEndDateStr = self.baseEndDateStr

        windFactorNameStr = 'adjfactor'
        colNameStr = 'adjfactor'
        dbUpdateDateFiledStr = 'adjfactor_update_date'
        dbUpdateDateFieldAliasStr = ' adjust factor '

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.writeBaseDataToMongoDB(baseStartDateStr, baseEndDateStr, windFactorNameStr, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr)


    def writeBaseIndicesCloseFromWindToMongoDB(self):
        baseStartDateStr = self.baseStartDateStr
        baseEndDateStr = self.baseEndDateStr

        windFactorNameStr = 'close'
        colNameStr = 'indices_close'
        dbUpdateDateFiledStr = 'index_close_update_date'
        dbUpdateDateFieldAliasStr = ' index close '

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.writeBaseDataToMongoDB(baseStartDateStr, baseEndDateStr, windFactorNameStr, colNameStr,
                                   dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr)

    def updateStocksCloseToMongoDB(self):

        windFactorNameStr = 'close'
        colNameStr = 'daily_close'
        dbUpdateDateFiledStr = 'stock_close_update_date'
        dbUpdateDateFieldAliasStr = ' daily close '

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.updateDataToMongoDB(windFactorNameStr, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr)

    def updateAdjFactorToMongoDB(self):
        windFactorNameStr = 'adjfactor'
        colNameStr = 'adjfactor'
        dbUpdateDateFiledStr = 'adjfactor_update_date'
        dbUpdateDateFieldAliasStr = ' adjust factor '

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.updateDataToMongoDB(windFactorNameStr, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr)


    def updateIndicesCloseToMongoDB(self):

        windFactorNameStr = 'close'
        colNameStr = 'indices_close'
        dbUpdateDateFiledStr = 'index_close_update_date'
        dbUpdateDateFieldAliasStr = ' index close '

        funcObj = StockDailyFactorGeneralFunctions()

        funcObj.updateDataToMongoDB(windFactorNameStr, colNameStr, dbUpdateDateFiledStr, dbUpdateDateFieldAliasStr)


    def writeFromMongoDBtoTxt(self, startDateStr, colStocksName, colIndicesName, txtFileName):

        ###############################################################################################################
        funcObj = StockDailyFactorGeneralFunctions()
        DBIP = funcObj.DBIP
        DBPort = funcObj.DBPort
        DBName = funcObj.DBName


        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_dailyclose = db[colStocksName]
        collection_indicesclose = db[colIndicesName]

        dateList = []
        totalStockCloseList = []

        dateListRes = collection_dailyclose.find({'StockID': '000001.SZ'})
        dateListDict = dateListRes[0]

        locateDateStrList =[]

        matchDateStr = datetime.datetime.strptime(startDateStr, '%Y-%m-%d').strftime('%Y%m')

        for (k, v) in dateListDict.items():
            try:
                dateStr = datetime.datetime.strptime(k, '%Y%m%d').strftime('%Y%m%d')
            except:
                dateStr = k

            try:
                locateDateStr = datetime.datetime.strptime(k, '%Y%m%d').strftime('%Y%m')
            except:
                locateDateStr = k

            dateList.append(dateStr)
            locateDateStrList.append(locateDateStr)

        dateList = dateList[2:]
        locateDateStrList = locateDateStrList[2:]

        dateNdx = locateDateStrList.index(matchDateStr)
        dateList = dateList[dateNdx:]
        dateList.insert(0, 'close')
        totalStockCloseList.append(dateList)
        totalSelectedDateNum = len(dateList)

        #######################################
        # Retrieve stock close data
        #######################################
        for stockIdRes in collection_dailyclose.find({}, {'StockID': 1}):
            stockIdStr = stockIdRes['StockID']
            for stockDailyCloseResDict in collection_dailyclose.find({'StockID': stockIdStr}):
                oneStockCloseList = []
                oneStockCloseList.append(stockIdStr)
                for i in range(1, totalSelectedDateNum):
                    try:
                        oneStockCloseList.append(stockDailyCloseResDict[dateList[i]])
                    except:
                        oneStockCloseList.append('0')
                totalStockCloseList.append(oneStockCloseList)


        #######################################
        # Retrieve indices close data
        #######################################
        stockListSize = len(totalStockCloseList)
        for stockIdRes in collection_indicesclose.find({}, {'StockID': 1}):
            stockIdStr = stockIdRes['StockID']
            for indexDailyCloseResDict in collection_indicesclose.find({'StockID': stockIdStr}):
                oneStockCloseList = []
                oneStockCloseList.append(stockIdStr)
                for i in range(1, totalSelectedDateNum):
                    try:
                        oneStockCloseList.append(indexDailyCloseResDict[dateList[i]])
                    except:
                        oneStockCloseList.append('0')
                totalStockCloseList.append(oneStockCloseList)

        print('Finished retrieving close data from MongoDB...')


        f = open(txtFileName, 'w')

        txtFileRowNum = len(oneStockCloseList)
        txtFileColNum = len(totalStockCloseList)

        f.write(str(txtFileRowNum - 1) + '\t' + str(txtFileColNum - 1) + '\t')

        for firstRowLen in range(txtFileColNum - 1):
            f.write('\t')
        f.write('\n')
        print('Finished writing text file header...')

        for rowIter in range(txtFileRowNum):
            for colIter in range(txtFileColNum):
                try:
                    cellData = totalStockCloseList[colIter][rowIter]
                    if math.isnan(float(cellData)):
                        cellData = 0
                    f.write(str(cellData) + '\t')
                except ValueError:
                    cellData = str(totalStockCloseList[colIter][rowIter])
                    f.write(cellData + '\t')

            f.write('\n')
            print('Wrote row number # ' + str(rowIter) + ' to the text file....')

        f.close()
        print('Finished writing the whole text file!')







obj = DailyClose()
#obj.writeBaseStocksCloseFromWindToMongoDB()
#obj.updateStocksCloseToMongoDB()

#obj.writeBaseIndicesCloseFromWindToMongoDB()
#obj.updateIndicesCloseToMongoDB()

#obj.writeBaseAdjFactorFromWindToMongoDB()
#obj.updateAdjFactorToMongoDB()


### to write retrace file, run the following two lines
backtraceStartDateStr = '2014-11-30'

obj.writeFromMongoDBtoTxt(backtraceStartDateStr, 'daily_close', 'indices_close', 'close.txt')